﻿//-----------------------------------------------------------------------------
// <copyright file="DependencyObjectExtensions.cs" company="William E. Kempf">
//     Copyright (c) William E. Kempf.
// </copyright>
//-----------------------------------------------------------------------------

namespace Onyx.Windows
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Windows;
    using System.Windows.Media;
    using Onyx.Reflection;

    /// <summary>
    /// Provides extension methods for <see cref="DependencyObject"/> types.
    /// </summary>
    public static class DependencyObjectExtensions
    {
        /// <summary>
        /// Gets the logical root of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the logical root for.</param>
        /// <returns>
        /// The logical root of the <see cref="DependencyObject"/>.
        /// </returns>
        public static DependencyObject GetLogicalRoot(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return self.GetLogicalAncestors().LastOrDefault() ?? self;
        }

        /// <summary>
        /// Gets the visual root of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the visual root for.</param>
        /// <returns>
        /// The visual root of the <see cref="DependencyObject"/>.
        /// </returns>
        public static DependencyObject GetVisualRoot(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return self.GetVisualAncestors().LastOrDefault() ?? self;
        }

        /// <summary>
        /// Gets the logical parent of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the logical parent for.</param>
        /// <returns>
        /// The logical parent of the <see cref="DependencyObject"/>.
        /// </returns>
        public static DependencyObject GetLogicalParent(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return LogicalTreeHelper.GetParent(self);
        }

        /// <summary>
        /// Gets the visual parent of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the visual parent for.</param>
        /// <returns>
        /// The visual parent of the <see cref="DependencyObject"/>.
        /// </returns>
        public static DependencyObject GetVisualParent(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return VisualTreeHelper.GetParent(self);
        }

        /// <summary>
        /// Gets the logical ancestors of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the logical ancestors for.</param>
        /// <returns>
        /// The logical ancestors of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetLogicalAncestors(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            DependencyObject ancestor = GetLogicalParent(self);
            while (ancestor != null)
            {
                yield return ancestor;
                ancestor = GetLogicalParent(ancestor);
            }
        }

        /// <summary>
        /// Gets the visual ancestors of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the visual ancestors for.</param>
        /// <returns>
        /// The visual ancestors of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetVisualAncestors(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            DependencyObject ancestor = GetVisualParent(self);
            while (ancestor != null)
            {
                yield return ancestor;
                ancestor = GetVisualParent(ancestor);
            }
        }

        /// <summary>
        /// Gets the logical children of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the logical children for.</param>
        /// <returns>
        /// The logical children of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetLogicalChildren(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return LogicalTreeHelper.GetChildren(self).OfType<DependencyObject>();
        }

        /// <summary>
        /// Gets the visual children of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the visual children for.</param>
        /// <returns>
        /// The visual children of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetVisualChildren(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            int count = VisualTreeHelper.GetChildrenCount(self);
            for (int i = 0; i < count; ++i)
            {
                yield return VisualTreeHelper.GetChild(self, i);
            }
        }

        /// <summary>
        /// Gets the logical descendents of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the logical descendents for.</param>
        /// <returns>
        /// The logical descendents of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetLogicalDescendents(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return GetDescendents(self, GetLogicalChildren, true);
        }

        /// <summary>
        /// Gets the logical descendents of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the logical descendents for.</param>
        /// <param name="depthFirst">If set to <see langword="true"/> then a depth first search for descendents will be
        /// made; otherwise a breadth first search will be made.</param>
        /// <returns>
        /// The logical descendents of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetLogicalDescendents(this DependencyObject self, bool depthFirst)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return GetDescendents(self, GetLogicalChildren, depthFirst);
        }

        /// <summary>
        /// Gets the visual descendents of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the visual descendents for.</param>
        /// <returns>
        /// The visual descendents of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetVisualDescendents(this DependencyObject self)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return GetDescendents(self, GetVisualChildren, true);
        }

        /// <summary>
        /// Gets the visual descendents of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="self">The <see cref="DependencyObject"/> to get the visual descendents for.</param>
        /// <param name="depthFirst">If set to <see langword="true"/> then a depth first search for descendents will be
        /// made; otherwise a breadth first search will be made.</param>
        /// <returns>
        /// The visual descendents of the <see cref="DependencyObject"/>.
        /// </returns>
        public static IEnumerable<DependencyObject> GetVisualDescendents(this DependencyObject self, bool depthFirst)
        {
            if (self == null)
            {
                throw new ArgumentNullException(Reflect.GetField(() => self).Name);
            }

            return GetDescendents(self, GetVisualChildren, depthFirst);
        }

        /// <summary>
        /// Gets the descendents of the <see cref="DependencyObject"/>.
        /// </summary>
        /// <param name="reference">The <see cref="DependencyObject"/> to get the descendents for.</param>
        /// <param name="getChildren">A delegate used to get the children of a <see cref="DependencyObject"/>.</param>
        /// <param name="depthFirst">If set to <see langword="true"/> then a depth first search for descendents will be
        /// made; otherwise a breadth first search will be made.</param>
        /// <returns>
        /// The descendents of the <see cref="DependencyObject"/>.
        /// </returns>
        private static IEnumerable<DependencyObject> GetDescendents(DependencyObject reference, Func<DependencyObject, IEnumerable<DependencyObject>> getChildren, bool depthFirst)
        {
            if (depthFirst)
            {
                Stack<DependencyObject> stack = new Stack<DependencyObject>();
                stack.Push(reference);
                while (stack.Count > 0)
                {
                    DependencyObject descendent = stack.Pop();
                    if (descendent != reference)
                    {
                        yield return descendent;
                    }

                    foreach (DependencyObject d in getChildren(descendent).Reverse())
                    {
                        stack.Push(d);
                    }
                }
            }
            else
            {
                Queue<DependencyObject> queue = new Queue<DependencyObject>();
                queue.Enqueue(reference);
                while (queue.Count > 0)
                {
                    DependencyObject descendent = queue.Dequeue();
                    if (descendent != reference)
                    {
                        yield return descendent;
                    }

                    foreach (DependencyObject d in getChildren(descendent))
                    {
                        queue.Enqueue(d);
                    }
                }
            }
        }
    }
}
